import oneflow as torch

def get_transformer_decoder_mask(targets):
    batch_size, steps = targets.size()
    seq_mask = torch.ones([batch_size, steps, steps], device=targets.device)
    seq_mask = torch.tril(seq_mask).to(torch.int8)
    return seq_mask


def get_streaming_frame_mask(tensor, length, left_context, right_context):
    B, T, _ = tensor.size()

    # if left_context + 1 > tuple:
    #     left_context = T - 1

    if left_context >= 0:
        left_mask = torch.tril(torch.ones((T, T), dtype=torch.uint8), diagonal=-(left_context + 1))
    else:
        left_mask = torch.zeros((T, T), dtype=torch.uint8)

    if right_context >= 0:
        right_mask = torch.triu(torch.ones((T, T), dtype=torch.uint8), diagonal=right_context + 1)
    else:
        right_mask = torch.zeros((T, T), dtype=torch.uint8)

    mask = left_mask + right_mask
    mask = mask.unsqueeze(0).repeat(B, 1, 1) < 1
    mask = mask.to(tensor.device)

    for b in range(tensor.size(0)):
        mask[b, length[b].item():] = True
        mask[b, :, length[b].item():] = False

    return mask.to(tensor.device)